package com.md.client.helper;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;

public class BodyCachingResponseWrapper extends HttpServletResponseWrapper {
  private final ByteArrayOutputStream outputStream;
  private final PrintWriter writer;

  public BodyCachingResponseWrapper(final HttpServletResponse response) {
    super(response);
    outputStream = new ByteArrayOutputStream();
    writer = new PrintWriter(outputStream);
  }

  @Override
  public ServletOutputStream getOutputStream() {
    return new ServletOutputStream() {
      @Override
      public boolean isReady() {
        return true;
      }

      @Override
      public void setWriteListener(final WriteListener writeListener) {

      }

      @Override
      public void write(final int b) {
        outputStream.write(b);
      }
    };
  }

  @Override
  public PrintWriter getWriter() {
    return writer;
  }

  public byte[] getCapturedData() {
    return outputStream.toByteArray();
  }

  public void resetResponse(final HttpServletResponse response) throws IOException {
    // 设置响应的状态码（保持不变）
    response.setStatus(getStatus());

    // 设置所有响应头
    for (String headerName : getHeaderNames()) {
      response.setHeader(headerName, getHeader(headerName));
    }

    // 获取字符编码并设置
    response.setCharacterEncoding("utf-8");
    // 将捕获的数据写回到客户端
    response.getOutputStream().write(outputStream.toByteArray());
    response.getOutputStream().flush();
  }
}
